"""
图片竖直合并

* 将多页pdf合并为一页
"""

import os
import cv2
import numpy as np

dirpath = 'temp'

output_fname = 'output.png'
tmp_fname = 'tmp.png'


def add_alpha_channel(img):
    """ 为jpg图像添加alpha通道 """

    b_channel, g_channel, r_channel = cv2.split(img)  # 剥离jpg图像通道
    alpha_channel = np.ones(b_channel.shape, dtype=b_channel.dtype) * 255  # 创建Alpha通道

    img_new = cv2.merge((b_channel, g_channel, r_channel, alpha_channel))  # 融合通道
    return img_new


def check_is_need_to_compress_size(img, base_shape_2, force_resize=False):
    need_compress = False
    if base_shape_2:
        shape = img.shape
        if force_resize or (shape[0] > base_shape_2[0] or shape[1] > base_shape_2[1]) or force_resize == '__all__':
            img = cv2.resize(img, base_shape_2)
            print('check_is_need_to_compress_size')
            need_compress = True
    return img, need_compress


def get_one_image(img_list, uniform_size=None, add_interval=True, force_resize=False, background_color=0, interval_color=200):
    """
    将多个图片合并为一个

    :param img_list: 图片列表
    :param uniform_size: 对过大的图片进行尺寸压缩
    :param add_interval: 增加图片间隔
    :param force_resize: 强制统一图片最大尺寸, 满足check_is_need_to_compress_size时才会运行. 如果为"__all__", 则统一所有图片.
    :param background_color: 背景颜色
    :param interval_color: 间隔颜色
    :return: concat后的图片
    """
    # img_list_0 = img_list.copy()
    # img_list = [img_list[i] for i in range(len(img_list)) if not isinstance(img_list[i], np.ndarray) and img_list[i] not in [None, '']]       # 提出空图片

    _img_list = []
    for i in range(len(img_list)):
        img_i = img_list[i]
        if isinstance(img_list[i], np.ndarray) or img_i not in [None, ''] :
            _img_list.append(img_i)
    img_list = _img_list

    # --- 理论最长
    max_width = 0

    for img in img_list:
        if isinstance(img, str):
            img = cv2.imread(img)
        print(img.shape)
        _, need_compress = check_is_need_to_compress_size(img, base_shape_2=uniform_size, force_resize=force_resize)
        if need_compress:
            max_width = uniform_size[0]
        else:
            if max_width < img.shape[1]:
                max_width = img.shape[1]

        if img.shape[1] > max_width or (uniform_size is not None and img.shape[0] > uniform_size[0]):
            max_width = img.shape[1]
    if uniform_size and uniform_size[0] < max_width:
        max_width = uniform_size[0]

    # --- 预处理
    interval_height = 20
    img_ls = []
    for img in img_list:
        if isinstance(img, str):
            assert os.path.exists(img), f'图片{img}不存在!'
            img = cv2.imread(img)
        if uniform_size:
            img, need_compress = check_is_need_to_compress_size(img, uniform_size, force_resize=force_resize)
        img_ls.append(img)
        if uniform_size and add_interval:
            interval_img = np.ones((interval_height, max_width, 3)) * interval_color  # 间隔图片
            img_ls.append(interval_img)
    # img_ls.pop(-1)
    for i in img_ls:
        print(i.shape)
    img_list = img_ls

    total_height = 1  # padding
    for img in img_list:
        if isinstance(img, str):
            img = cv2.imread(img)
        if img.shape[1] > max_width:
            max_width = img.shape[1]
        total_height += img.shape[0]

    # create a new array with a size large enough to contain all the images
    final_image = np.ones((total_height, max_width, 3), dtype=np.uint8) * background_color

    current_y = 0  # keep track of where your current image was last placed in the y coordinate
    for image in img_list:
        # add an image to the final array and increment the y coordinate

        tmp_img = np.zeros((image.shape[0], max_width - image.shape[1], 3))
        image = np.hstack((image, tmp_img))
        final_image[current_y:current_y + image.shape[0], :, :] = image
        current_y += image.shape[0]
    return final_image


def compose_two_image_with_hstack(left, right):
    img1 = cv2.imread(left)
    img2 = cv2.imread(right)
    image = np.hstack((img1, img2))
    cv2.imwrite("left_right.png", image)


def compose_two_image_with_vstack(up, down):
    img1 = cv2.imread(up)
    img2 = cv2.imread(down)
    image = np.vstack((img1, img2))
    cv2.imwrite(tmp_fname, image)


if __name__ == '__main__':
    imgs = [os.path.join(dirpath, img) for img in os.listdir(dirpath)]

    image = get_one_image(imgs)
    print(image)
    cv2.imshow("test", image)
    cv2.waitKey(0)

    base_shape_2 = [600, 300]
    image = get_one_image(imgs, uniform_size=base_shape_2, force_resize='__all__', add_interval=True, background_color=200)
    cv2.imshow("test", image)
    cv2.waitKey(0)

    assert len(imgs) >= 2, '图片数量要大于两张!'
    compose_two_image_with_vstack(imgs[0], imgs[1])

    if len(imgs) > 2:
        for img in imgs[2:]:
            compose_two_image_with_vstack(tmp_fname, img)

    if os.path.exists(output_fname):
        os.remove(output_fname)
    os.rename(tmp_fname, output_fname)

    if os.path.exists(tmp_fname):
        os.remove(tmp_fname)

    os.startfile(output_fname)
    print('--- 图片输出路径: ', os.path.abspath(output_fname))

